package com.chatplus.application.listener;

import cn.bugstack.openai.session.Configuration;
import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
import com.chatplus.application.aiprocessor.platform.image.sd.SdImageProcessor;
import com.chatplus.application.common.logging.SouthernQuietLogger;
import com.chatplus.application.common.logging.SouthernQuietLoggerFactory;
import com.chatplus.application.common.util.OkHttpClientUtil;
import com.chatplus.application.common.util.PlusJsonUtils;
import com.chatplus.application.domain.dto.ApiKeyDto;
import com.chatplus.application.domain.dto.SdJobResult;
import com.chatplus.application.domain.entity.draw.SdJobEntity;
import com.chatplus.application.domain.notification.SdLocalJobNotification;
import com.chatplus.application.enumeration.AiPlatformEnum;
import com.chatplus.application.service.draw.SdJobService;
import com.chatplus.application.web.idempotent.annotation.Concurrent;
import com.chatplus.application.web.notification.NotificationListener;
import okhttp3.*;
import org.apache.commons.lang3.StringUtils;
import org.springframework.stereotype.Component;

import java.io.IOException;

/**
 * Sd 任务监听
 */
@Component
public class SdLocalJobNotificationListener {

    private final SdImageProcessor sdImageProcessor;
    private final SdJobService sdJobService;
    private final SdJobProcess sdJobProcess;

    public SdLocalJobNotificationListener(SdImageProcessor sdImageProcessor,
                                          SdJobService sdJobService,
                                          SdJobProcess sdJobProcess) {
        this.sdImageProcessor = sdImageProcessor;
        this.sdJobService = sdJobService;
        this.sdJobProcess = sdJobProcess;
    }

    private static final SouthernQuietLogger LOGGER = SouthernQuietLoggerFactory.getLogger(SdLocalJobNotificationListener.class);

    @NotificationListener(notification = SdLocalJobNotification.class)
    @Concurrent(value = "SdLocalJobNotificationListener-handleSdLocalJobNotification")
    public void handleSdLocalJobNotification(SdLocalJobNotification notification) {
        ApiKeyDto apiKeyDto = notification.getApiKeyDto();
        SdJobEntity sdJobEntity = sdJobService.getById(notification.getSdJobId());
        if (sdJobEntity == null) {
            return;
        }
        OkHttpClient okHttpClient = OkHttpClientUtil.getOkHttpClient(apiKeyDto.getProxyUrl(), true);
        Request.Builder builder = new Request.Builder()
                .addHeader("Content-Type", Configuration.APPLICATION_JSON)
                .addHeader("Accept", Configuration.APPLICATION_JSON);
        builder.url(apiKeyDto.getUrl() + "/sdapi/v1/txt2img")
                .post(RequestBody.create(PlusJsonUtils.toJsonString(notification.getRequest()), MediaType.parse(Configuration.APPLICATION_JSON)));
        Request request = builder.build();
        Call call = okHttpClient.newCall(request);
        sdJobProcess.checkTaskProgress(sdJobEntity.getId(), apiKeyDto);
        try (Response response = call.execute()) {
            if (response.isSuccessful() && response.body() != null) {
                String body = response.body().string();
                SdJobResult result = PlusJsonUtils.parseObject(body, SdJobResult.class);
                // ----------- 检查返回是否成功 ------------
                if (result == null || !result.isSuccess(AiPlatformEnum.SD_LOCAL)) {
                    LOGGER.message("SD Local请求返回为空").context("response", response).error();
                    sdJobProcess.handlerError(sdJobEntity, AiPlatformEnum.SD_LOCAL);
                    return;
                }
                // ----------上传base64图片------------
                String base64Image = result.getImages().getFirst();
                String url = sdImageProcessor.saveToOss(null, base64Image);
                if (StringUtils.isEmpty(url)) {
                    LOGGER.message("SD Base64图片上传失败").context("response", body).error();
                    sdJobProcess.handlerError(sdJobEntity, AiPlatformEnum.SD_LOCAL);
                    return;
                }
                sdJobEntity.setImgUrl(url);
                sdJobEntity.setProgress(100);
            } else {
                LOGGER.message("SD Local图片请求失败").context("response", response).error();
                sdJobProcess.handlerError(sdJobEntity, AiPlatformEnum.SD_LOCAL);
            }
        } catch (IOException e) {
            LOGGER.message("SD Local图片请求失败").context("exception", e).error();
            sdJobProcess.handlerError(sdJobEntity, AiPlatformEnum.SD_LOCAL);
        } finally {
            LambdaUpdateWrapper<SdJobEntity> updateWrapper = new LambdaUpdateWrapper<>();
            updateWrapper.eq(SdJobEntity::getId, sdJobEntity.getId());
            updateWrapper.set(SdJobEntity::getProgress, sdJobEntity.getProgress());
            updateWrapper.set(StringUtils.isNotEmpty(sdJobEntity.getImgUrl()), SdJobEntity::getImgUrl, sdJobEntity.getImgUrl());
            sdJobService.update(updateWrapper);
            sdImageProcessor.notifyUpdateTask(sdJobEntity.getId());
        }
    }


}
